import mss
import mss.tools
import cv2
import numpy as np
import time
import pyautogui
import threading
import keyboard

global_stop_recording = False
snipaste_screen = False

# 新建一个主线程作为显示屏幕是否正录制的GUI(pyautogui)
def displat_recording():
    pyautogui.alert("正在录制屏幕...\n按 'q' 键停止")
    # 跳出按钮
    if pyautogui.confirm("是否确认停止录制？") == "OK":
        global global_stop_recording
        global_stop_recording = True

def set_snipaste_screen():
    """
    设置截图状态为True
    """
    global snipaste_screen
    snipaste_screen = True

# 将获取屏幕视频流的代码放到一个线程中
def record_screenstream():
    # 屏幕录制参数
    output_file = "screen_recording.avi"  # 输出文件名
    fps = 15  # 帧率（根据性能调整）

    # 设置视频编码器（根据格式选择，mp4常用XVID）
    fourcc = cv2.VideoWriter_fourcc(*"XVID")
    # fourcc = cv2.VideoWriter_fourcc(*"mp4v")

    print("开始录制，按 'q' 键停止...")
    with mss.mss() as sct:
        monitor = sct.monitors[1]  # 1表示主屏幕（0是所有屏幕合并）
        width, height = monitor["width"], monitor["height"]
        out = cv2.VideoWriter(output_file, fourcc, fps, (width, height))
        count = 0
        while not global_stop_recording:
            # 截取屏幕
            sct_img = sct.grab(monitor)
            # 转换为OpenCV格式（BGR）
            frame = np.array(sct_img)
            # frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2BGR)  # mss默认是RGBA，转RGB更贴近显示器的渲染方式
            frame = cv2.cvtColor(frame, cv2.COLOR_RGBA2RGB)  # 不用转换颜色吧
            # 写入视频，并确保写入成功
            if out.write(frame) == False:
                print("写入视频失败")
                break
            if keyboard.is_pressed("s"):
                print("截图键被按下")
                count += 1
                # 保存截图
                cv2.imwrite(f"../screen_frame/screenshot_{count}.png", frame)
                print(f"截图已保存为：screenshot_{count}.png")
            if global_stop_recording: # 检查是否需要停止录制
                break
        out.release()
        cv2.destroyAllWindows()
        print(f"录制完成，文件保存为：{output_file}")

if __name__ == "__main__":
    # 用Tread新建一个线程来显示正在录制的GUI
    recording_thread = threading.Thread(target=displat_recording)
    recording_thread.start()
    # 用Tread新建一个线程来截图
    snipaste_thread = threading.Thread(target=set_snipaste_screen)
    snipaste_thread.start()
    stream_thread = threading.Thread(target=record_screenstream)
    # 注册为守护线程，主线程结束时，子线程也会结束
    stream_thread.daemon = True # 设置线程属性要在启动之前(untimeError: cannot set daemon status of active thread)
    stream_thread.start()

